youtube
Are you exploring JAX for the first time and feeling overwhelmed by terms like "functional purity," "explicit state," and "jit"? You aren't alone.
Moving from traditional object-oriented machine learning (like PyTorch or TensorFlow) to JAX requires a shift in your mental model. In this video, we break down exactly what it means to program in a "functional" world and how that shift unlocks blazing-fast performance and scalability for your models.
We cover the constraints of high performance ML: strict rules on how you handle variables, random numbers, and side effects. By the end, you'll understand why JAX asks you to be explicit about state and how XLA and JIT compilation optimize your code.
Resources:
Documentation on Thinking in JAX→
Functional Purity →
JIT →
Pseudorandom numbers →
Check out Flax (Neural Networks on JAX) →
Chapters:
0:00 Introduction to JAX concepts
0:30 What is JAX?
1:03 Functional Purity in JAX
2:15 Non-Modifiable Arrays
3:12 Explicit State Management
4:47 Pseudo Random Number Generation (PRNG)
6:30 Just-In-Time (JIT) Compilation
8:06 JIT Limitations & Control Flow
9:11 Summary & Conclusion
Subscribe to Google for Developers →
Speaker: Yufeng Guo,
Products Mentioned: Keras, Gemma, JAX
|
2026/01/15
youtube
